#include <Arduino.h>
#include <ESPAsyncWebServer.h>
#include <WiFiUdp.h>
#include "StreamString.h"
#include "ESPAsyncHTTPUpdateServer.h"

static const char serverIndex[] PROGMEM =
    R"(<html><body><form method='POST' action='' enctype='multipart/form-data'>
                  <input type='file' name='update'>
                  <input type='submit' value='Update'>
               </form>
         </body></html>)";
static const char successResponse[] PROGMEM =
    "<META http-equiv=\"refresh\" content=\"15;URL=/\">Update Success! Rebooting...\n";

ESPAsyncHTTPUpdateServer::ESPAsyncHTTPUpdateServer(bool serial_debug)
{
    _serial_output = serial_debug;
    _server = NULL;
    _username = emptyString;
    _password = emptyString;
    _authenticated = false;
}

ESPAsyncHTTPUpdateServer::~ESPAsyncHTTPUpdateServer()
{
}

void ESPAsyncHTTPUpdateServer::setup(AsyncWebServer *server, const String &path, const String &username, const String &password)
{
    _server = server;
    _username = username;
    _password = password;

    // handler for the /update form page
    _server->on(path.c_str(), HTTP_GET, [&](AsyncWebServerRequest *request) {
        if (_username != emptyString && _password != emptyString && !request->authenticate(_username.c_str(), _password.c_str()))
            return request->requestAuthentication();
        request->send_P(200, PSTR("text/html"), serverIndex);
    });

    // handler for the /update form POST (once file upload finishes)
    _server->on(path.c_str(), HTTP_POST,
                [&](AsyncWebServerRequest *request) {
                    if (!_authenticated)
                        return request->requestAuthentication();
                    if (Update.hasError())
                    {
                        request->send(200, F("text/html"), String(F("Update error: ")) + _updaterError);
                    }
                    else
                    {
                        request->client()->setNoDelay(true);
                        AsyncWebServerResponse *response = request->beginResponse(200, PSTR("text/html"), successResponse);
                        response->addHeader("Connection", "close");
                        request->send(response);
                        delay(100);
                        request->client()->stop();

                        delay(5000);
                        ESP.reset();
                        delay(2000);
                        // request->send_P(200, PSTR("text/html"), successResponse);
                    }
                },
                [&](AsyncWebServerRequest *request, const String &filename, size_t index, uint8_t *data, size_t len, bool final) {
                    if (!index)
                    {
                        _updaterError = String();
                        if (_serial_output)
                            Serial.setDebugOutput(true);

                        _authenticated = (_username == emptyString || _password == emptyString || request->authenticate(_username.c_str(), _password.c_str()));
                        if (!_authenticated)
                        {
                            if (_serial_output)
                                Serial.printf("Unauthenticated Update\n");
                            return;
                        }

                        WiFiUDP::stopAll();
                        if (_serial_output)
                            Serial.printf("Update: %s\n", filename.c_str());

                        Update.runAsync(true);
                        uint32_t maxSketchSpace = (ESP.getFreeSketchSpace() - 0x1000) & 0xFFFFF000;
                        if (!Update.begin(maxSketchSpace))
                        { //start with max available size
                            _setUpdaterError();
                        }
                    }
                    if (_authenticated && !_updaterError.length() && !Update.hasError())
                    {
                        if (_serial_output)
                            Serial.printf(".");
                        if (Update.write(data, len) != len)
                        {
                            _setUpdaterError();
                        }
                    }
                    if (_authenticated && final && !_updaterError.length())
                    {
                        if (Update.end(true))
                        { //true to set the size to the current progress
                            if (_serial_output)
                                Serial.printf("Update Success: %u\nRebooting...\n", index + len);
                        }
                        else
                        {
                            _setUpdaterError();
                        }
                        if (_serial_output)
                            Serial.setDebugOutput(false);
                    }
                    // if (_authenticated)
                    // {
                    //     Update.end();
                    //     if (_serial_output)
                    //         Serial.println("Update was aborted");
                    // }
                });
}

void ESPAsyncHTTPUpdateServer::_setUpdaterError()
{
    if (_serial_output)
        Update.printError(Serial);
    StreamString str;
    Update.printError(str);
    _updaterError = str.c_str();
}
